package cipher

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"os"
	"strings"
)

func checkError(err error) {
	if nil != err {
		panic(err)
	}
}

func GetKey(key string) string {
	array := strings.Split(key, "\n")
	var result string
	for i := 1; i < len(array)-2; i++ {
		result += array[i]
	}
	return result
}

func SaveRSAKey(key_file string, key_byte []byte) {
	//创建文件保存私钥
	privateFile, err := os.Create(key_file)
	if err != nil {
		panic(err)
	}
	defer privateFile.Close()
	//构建一个pem.Block结构体对象
	privateBlock := pem.Block{Type: "RSA Private Key", Bytes: key_byte}
	//将数据保存到文件
	pem.Encode(privateFile, &privateBlock)
}

//生成RSA私钥和公钥，保存到文件中
func GenerateRSAKey(bits int, private_file string, public_file string) {
	//GenerateKey函数使用随机数据生成器random生成一对具有指定字位数的RSA密钥
	//Reader是一个全局、共享的密码用强随机数生成器
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		panic(err)
	}
	//保存私钥
	//通过x509标准将得到的ras私钥序列化为ASN.1 的 DER编码字符串
	X509PrivateKey := x509.MarshalPKCS1PrivateKey(privateKey)
	SaveRSAKey(private_file, X509PrivateKey)

	//保存公钥
	//获取公钥的数据
	publicKey := privateKey.PublicKey
	//X509对公钥编码
	X509PublicKey, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		panic(err)
	}
	SaveRSAKey(public_file, X509PublicKey)
}

func ReadFile(path string) string {
	file, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer file.Close()
	info, _ := file.Stat()
	buf := make([]byte, info.Size())
	file.Read(buf)
	return string(buf)
}

func PemDecode(buf string) []byte {
	// pem解码
	block, _ := pem.Decode([]byte(buf))
	return block.Bytes
}

//RSA加密
func RSA_Encrypt(plainText string, key string) string {
	pem_key := PemDecode(key)
	//x509解码
	publicKeyInterface, err := x509.ParsePKIXPublicKey(pem_key)
	checkError(err)
	//类型断言
	publicKey := publicKeyInterface.(*rsa.PublicKey)
	//对明文进行加密
	cipherText, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, []byte(plainText))
	checkError(err)
	result := base64.StdEncoding.EncodeToString([]byte(cipherText))
	//返回密文
	return result
}

//RSA解密
func RSA_Decrypt(cipherText string, key string) string {
	pwd_srt, err := base64.StdEncoding.DecodeString(cipherText)
	checkError(err)
	pem_key := PemDecode(key)
	//X509解码
	privateKey, err := x509.ParsePKCS1PrivateKey(pem_key)
	checkError(err)
	//对密文进行解密
	plainText, _ := rsa.DecryptPKCS1v15(rand.Reader, privateKey, pwd_srt)
	//返回明文
	return string(plainText)
}

func GenerateRSAKeyPair() map[string]string {
	private_file := "rsa_private_key.pem"
	public_file := "rsa_public_key.pem"
	GenerateRSAKey(2048, private_file, public_file)
	public_key := ReadFile(public_file)
	private_key := ReadFile(private_file)
	os.Remove(private_file)
	os.Remove(public_file)
	result := map[string]string{
		"private_key": private_key,
		"public_key":  public_key,
	}
	return result
}

// func main() {
// 	private_file := "rsa_private_key.pem"
// 	public_file := "rsa_public_key.pem"
// 	//生成密钥对，保存到文件
// 	GenerateRSAKey(2048, private_file, public_file)
// 	public_key := ReadFile(public_file)
// 	private_key := ReadFile(private_file)
// 	message := "芝麻开门"
// 	//加密
// 	cipherText := RSA_Encrypt(message, public_key)
// 	sEnc := base64.StdEncoding.EncodeToString([]byte(cipherText))
// 	fmt.Println("======sEnc====")
// 	fmt.Println(sEnc)
// 	fmt.Println("==========")
// 	fmt.Println(public_key)
// 	sDec, _ := base64.StdEncoding.DecodeString(sEnc)
// 	//解密
// 	a3 := string(sDec)
// 	plainText := RSA_Decrypt(a3, private_key)
// 	fmt.Println(string(plainText))
// 	fmt.Println(private_key)
// 	fmt.Println(len(private_key))
// 	os.Remove(private_file)
// 	os.Remove(public_file)
// }
